#!/usr/bin/env python
# -*- coding:utf-8 -*-
# @FileName  :nginx_log_ddos_alert.py
# @Time      :2025/07/25 15:30:10
# @Author    :zhaotp
"""
脚本名称：nginx_log_ddos_alert.py

功能说明：
  实时监控Nginx JSON格式访问日志，统计每个公网IP对每个URL（去除参数）一分钟内的请求次数，
  超过30次则通过钉钉机器人发送告警。
  当同一IP触发100次告警且是国外IP时，自动加入iptables拦截。
  注意：国内IP不会被自动拦截，仅发送告警通知。

主要特性：
  - 支持多日志文件（/usr/local/nginx/logs/*.log）
  - 适配JSON格式日志（需包含 remote_ip、url、timestamp 字段）
  - 统计时自动去除URL参数，仅以问号前的路径为准
  - 排除指定内网网段（如 172.17.6.0/24, 172.17.8.0/24）
  - 排除指定URL前缀（如 /app-xxx,/web/app/xxx/）
  - 自动排除静态资源文件（如 .js、.css、.png、.jpg、.gif、.ico、.svg、.woff、.ttf 等）
  - 钉钉预警内容包含环境、服务器、IP、接口、次数、时间等关键信息
  - 支持自定义钉钉Webhook
  - 国外IP检测功能
  - 自动iptables拦截功能（触发100次且为国外IP）

依赖：
  - requests
  - ipaddress
  - geoip2 (用于国外IP检测)

用法：
  python3 nginx_log_ddos_alert.py

自定义说明：
  - 如需排除更多网段，请在 EXCLUDE_NETS 中添加
  - 如需排除更多接口前缀，请在 EXCLUDE_URL_PREFIXES 中添加
  - 如需排除更多静态资源类型，请在 EXCLUDE_STATIC_EXTENSIONS 中添加
  - 统计逻辑已自动去除URL参数，聚合同一接口的不同参数请求

"""
import glob
import json
import time
import subprocess
from datetime import datetime, timedelta
from collections import defaultdict
import requests
import logging

# 先配置日志
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s [%(levelname)s] %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
)

# 尝试导入geoip2库，如果失败则使用备用方案
try:
    import geoip2.database
    import geoip2.errors
    GEOIP2_AVAILABLE = True
    logging.info("geoip2库导入成功，将使用精确的地理位置检测")
except ImportError:
    GEOIP2_AVAILABLE = False
    logging.warning("geoip2库未安装，将使用简化的中国IP网段检测")

# 配置
LOG_PATH = "/usr/local/nginx/logs/*.log"
THRESHOLD = 30  # 告警阈值
BLOCK_THRESHOLD = 100  # 拦截阈值
INTERVAL = 60  # 秒
DINGTALK_WEBHOOK = "https://oapi.dingtalk.com/robot/send?access_token=XXXXXXX" #钉钉api接口

# GeoIP2数据库配置
GEOIP2_DB_PATH = "/data/venv/GeoLite2-Country_20250815/GeoLite2-Country.mmdb"  # 默认路径，请根据实际情况修改

# 判断是否公网IP，且不在排除网段
import ipaddress
#排除指定网段
EXCLUDE_NETS = [
    ipaddress.ip_network("172.17.6.0/24"),
    ipaddress.ip_network("172.17.8.0/24"),
]

# 中国IP网段（简化版，实际使用时建议使用geoip2库）
CHINA_IP_RANGES = [
    "1.0.1.0/24", "1.0.2.0/23", "1.0.8.0/21", "1.0.32.0/19", "1.0.64.0/18",
    "1.0.128.0/17", "1.1.0.0/24", "1.1.2.0/23", "1.1.4.0/22", "1.1.8.0/21",
    "1.1.16.0/20", "1.1.32.0/19", "1.1.64.0/18", "1.1.128.0/17", "1.2.0.0/15",
    "1.4.0.0/14", "1.8.0.0/13", "1.16.0.0/12", "1.32.0.0/11", "1.64.0.0/10",
    "1.128.0.0/9", "2.0.0.0/7", "4.0.0.0/6", "8.0.0.0/5", "16.0.0.0/4",
    "32.0.0.0/3", "64.0.0.0/2", "128.0.0.0/1", "14.0.0.0/13", "14.16.0.0/12",
    "14.32.0.0/11", "14.64.0.0/13", "14.80.0.0/12", "14.96.0.0/11", "14.112.0.0/12",
    "14.128.0.0/11", "14.144.0.0/12", "14.160.0.0/11", "14.176.0.0/12", "14.192.0.0/11",
    "14.208.0.0/12", "14.224.0.0/11", "14.240.0.0/12", "27.0.0.0/8", "27.8.0.0/13",
    "27.16.0.0/12", "27.32.0.0/11", "27.48.0.0/13", "27.56.0.0/12", "27.64.0.0/11",
    "27.80.0.0/12", "27.96.0.0/11", "27.112.0.0/12", "27.128.0.0/11", "27.144.0.0/12",
    "27.160.0.0/11", "27.176.0.0/12", "27.192.0.0/11", "27.208.0.0/12", "27.224.0.0/11",
    "27.240.0.0/12", "36.0.0.0/7", "38.0.0.0/8", "39.0.0.0/8", "40.0.0.0/6",
    "42.0.0.0/7", "44.0.0.0/6", "46.0.0.0/7", "48.0.0.0/5", "52.0.0.0/6",
    "54.0.0.0/7", "56.0.0.0/5", "58.0.0.0/7", "60.0.0.0/6", "62.0.0.0/7",
    "64.0.0.0/2", "128.0.0.0/1", "101.0.0.0/8", "103.0.0.0/8", "106.0.0.0/7",
    "108.0.0.0/6", "110.0.0.0/7", "112.0.0.0/5", "114.0.0.0/7", "116.0.0.0/6",
    "118.0.0.0/7", "120.0.0.0/6", "122.0.0.0/7", "124.0.0.0/6", "126.0.0.0/7",
    "128.0.0.0/1", "202.0.0.0/7", "204.0.0.0/6", "206.0.0.0/7", "208.0.0.0/6",
    "210.0.0.0/7", "212.0.0.0/6", "214.0.0.0/7", "216.0.0.0/6", "218.0.0.0/7",
    "220.0.0.0/6", "222.0.0.0/7", "223.0.0.0/8", "223.64.0.0/10", "223.68.0.0/11",
    "223.70.0.0/12", "223.72.0.0/11", "223.74.0.0/12", "223.76.0.0/11", "223.78.0.0/12",
    "223.80.0.0/11", "223.82.0.0/12", "223.84.0.0/11", "223.86.0.0/12", "223.88.0.0/11",
    "223.90.0.0/12", "223.92.0.0/11", "223.94.0.0/12", "223.96.0.0/11", "223.98.0.0/12",
    "223.100.0.0/11", "223.102.0.0/12", "223.104.0.0/11", "223.106.0.0/12", "223.108.0.0/11",
    "223.110.0.0/12", "223.112.0.0/11", "223.114.0.0/12", "223.116.0.0/11", "223.118.0.0/12",
    "223.120.0.0/11", "223.122.0.0/12", "223.124.0.0/11", "223.126.0.0/12", "223.128.0.0/11",
    "223.130.0.0/12", "223.132.0.0/11", "223.134.0.0/12", "223.136.0.0/11", "223.138.0.0/12",
    "223.140.0.0/11", "223.142.0.0/12", "223.144.0.0/11", "223.146.0.0/12", "223.148.0.0/11",
    "223.150.0.0/12", "223.152.0.0/11", "223.154.0.0/12", "223.156.0.0/11", "223.158.0.0/12",
    "223.160.0.0/11", "223.162.0.0/12", "223.164.0.0/11", "223.166.0.0/12", "223.168.0.0/11",
    "223.170.0.0/12", "223.172.0.0/11", "223.174.0.0/12", "223.176.0.0/11", "223.178.0.0/12",
    "223.180.0.0/11", "223.182.0.0/12", "223.184.0.0/11", "223.186.0.0/12", "223.188.0.0/11",
    "223.190.0.0/12", "223.192.0.0/11", "223.194.0.0/12", "223.196.0.0/11", "223.198.0.0/12",
    "223.200.0.0/11", "223.202.0.0/12", "223.204.0.0/11", "223.206.0.0/12", "223.208.0.0/11",
    "223.210.0.0/12", "223.212.0.0/11", "223.214.0.0/12", "223.216.0.0/11", "223.218.0.0/12",
    "223.220.0.0/11", "223.222.0.0/12", "223.224.0.0/11", "223.226.0.0/12", "223.228.0.0/11",
    "223.230.0.0/12", "223.232.0.0/11", "223.234.0.0/12", "223.236.0.0/11", "223.238.0.0/12",
    "223.240.0.0/11", "223.242.0.0/12", "223.244.0.0/11", "223.246.0.0/12", "223.248.0.0/11",
    "223.250.0.0/12", "223.252.0.0/11", "223.254.0.0/12"
]

def is_public_ip(ip):
    try:
        ip_obj = ipaddress.ip_address(ip)
        # 排除指定网段
        for net in EXCLUDE_NETS:
            if ip_obj in net:
                return False
        return ip_obj.is_global
    except Exception:
        return False

def is_foreign_ip(ip):
    """判断是否为国外IP（非中国IP）"""
    try:
        ip_obj = ipaddress.ip_address(ip)
        
        # 如果geoip2库可用，使用精确的地理位置检测
        if GEOIP2_AVAILABLE:
            try:
                # 尝试打开GeoIP2数据库
                with geoip2.database.Reader(GEOIP2_DB_PATH) as reader:
                    response = reader.country(ip)
                    country_code = response.country.iso_code
                    
                    # 检查是否为中国
                    if country_code == 'CN':
                        logging.debug("IP %s 检测为中国IP (国家代码: %s)", ip, country_code)
                        return False
                    else:
                        logging.debug("IP %s 检测为国外IP (国家代码: %s)", ip, country_code)
                        return True
                        
            except FileNotFoundError:
                logging.warning("GeoIP2数据库文件未找到: %s，将使用备用检测方法", GEOIP2_DB_PATH)
            except Exception as e:
                logging.warning("GeoIP2检测失败: %s，将使用备用检测方法", e)
        
        # 备用方案：使用简化的中国IP网段检测
        logging.debug("使用备用IP检测方法")
        for china_net in CHINA_IP_RANGES:
            if ip_obj in ipaddress.ip_network(china_net):
                return False
        return True
        
    except Exception as e:
        logging.error("IP检测异常: %s, IP: %s", e, ip)
        return False

def init_geoip2_database():
    """初始化GeoIP2数据库"""
    if not GEOIP2_AVAILABLE:
        logging.warning("geoip2库未安装，跳过数据库初始化")
        return False
    
    try:
        # 检查数据库文件是否存在
        import os
        if os.path.exists(GEOIP2_DB_PATH):
            # 测试数据库是否可用
            with geoip2.database.Reader(GEOIP2_DB_PATH) as reader:
                # 测试一个已知的中国IP
                test_ip = "114.114.114.114"  # 114DNS
                response = reader.country(test_ip)
                logging.info("GeoIP2数据库初始化成功，测试IP %s 检测为: %s", 
                           test_ip, response.country.name or response.country.iso_code)
                return True
        else:
            logging.warning("GeoIP2数据库文件不存在: %s", GEOIP2_DB_PATH)
            logging.info("请下载GeoIP2数据库文件并放置在正确路径")
            return False
    except Exception as e:
        logging.error("GeoIP2数据库初始化失败: %s", e)
        return False

def block_ip_with_iptables(ip):
    """使用iptables拦截IP"""
    try:
        # 检查IP是否已经被拦截
        check_cmd = f"iptables -C INPUT -s {ip} -j DROP"
        result = subprocess.run(check_cmd, shell=True, capture_output=True)
        
        if result.returncode != 0:  # IP未被拦截，执行拦截
            # 在INPUT链的末尾插入规则，确保优先级最高
            block_cmd = f"iptables -A INPUT -s {ip} -j DROP"
            result = subprocess.run(block_cmd, shell=True, capture_output=True)
            
            if result.returncode == 0:
                logging.warning("成功拦截IP: %s (规则已插入到INPUT链开头)", ip)
                return True
            else:
                logging.error("拦截IP失败: %s, 错误: %s", ip, result.stderr.decode())
                return False
        else:
            logging.info("IP %s 已经被拦截", ip)
            return True
    except Exception as e:
        logging.exception("拦截IP时发生异常: %s", e)
        return False

#排除指定url
EXCLUDE_URL_PREFIXES = [
    "/gateway-xxx",
    "/app-xxx",
    "/web/app/xxx/"
]

# 排除静态资源文件扩展名
EXCLUDE_STATIC_EXTENSIONS = [
    '.js', '.css', '.png', '.jpg', '.jpeg', '.gif', '.bmp', '.ico', '.svg',
    '.woff', '.woff2', '.ttf', '.eot', '.otf', '.pdf', '.zip', '.rar',
    '.mp4', '.mp3', '.avi', '.mov', '.wmv', '.flv', '.webm',
    '.xml', '.json', '.txt', '.log', '.md', '.html', '.htm'
]

def is_static_resource(url):
    """判断是否为静态资源"""
    if not url:
        return False
    # 检查URL是否以静态资源扩展名结尾
    return any(url.lower().endswith(ext) for ext in EXCLUDE_STATIC_EXTENSIONS)

def parse_log_line(line):
    try:
        data = json.loads(line)
        ip = data.get("remote_ip")
        url = data.get("url")
        timestamp = data.get("timestamp")
        # 只保留问号前的url路径
        if url and "?" in url:
            url = url.split("?")[0]
        # 解析时间戳
        dt = datetime.strptime(timestamp[:19], "%Y-%m-%dT%H:%M:%S")
        return ip, dt, url
    except Exception:
        return None, None, None

def send_dingtalk_alert(ip, url, count, alert_time, is_blocked=False, is_foreign=False):
    """发送钉钉告警"""
    # 获取IP地理位置信息
    location_info = get_ip_location_info(ip)
    
    # 计算请求频率（每秒）
    requests_per_second = count / 60
    
    if is_blocked and is_foreign:
        # 模板1：国外IP拦截告警
        content = f"""🚨【DDOS攻击自动拦截告警】🚨

🚨 国外IP {ip} 触发{count}次攻击，已自动拦截！

🔴 告警级别：严重
🌍 攻击来源：国外IP
📍 IP地址：{ip}
🌐 地理位置：{location_info}
📊 攻击强度：{count}次/分钟 ({requests_per_second:.1f}次/秒)
🔗 攻击接口：{url}
🖥️ 目标服务器：127.0.0.1
⏰ 告警时间：{alert_time}
✅ 处理状态：已自动拦截
🎯 拦截阈值：{BLOCK_THRESHOLD}次/分钟
"""
    else:
        # 模板2：其他所有情况的告警
        content = f"""⚠️【异常访问监控告警】⚠️

⚠️ IP {ip} 触发{count}次请求，需要关注

🟡 告警级别：警告
📍 IP地址：{ip}
🌐 地理位置：{location_info}
📊 请求频率：{count}次/分钟 ({requests_per_second:.1f}次/秒)
🔗 访问接口：{url}
🖥️ 目标服务器：127.0.0.1
⏰ 告警时间：{alert_time}
🎯 触发阈值：{THRESHOLD}次/分钟
"""
    
    data = {
        "msgtype": "text",
        "text": {"content": content}
    }
    
    alert_type = "拦截告警" if is_blocked else "普通告警"
    logging.warning("%s：公网ip=%s 接口=%s 次数=%s 时间=%s", alert_type, ip, url, count, alert_time)
    
    try:
        resp = requests.post(DINGTALK_WEBHOOK, json=data)
        logging.info("钉钉预警发送结果: %s", resp.text)
    except Exception:
        logging.exception("钉钉预警发送失败")

def get_ip_location_info(ip):
    """获取IP地理位置信息"""
    try:
        if GEOIP2_AVAILABLE:
            try:
                with geoip2.database.Reader(GEOIP2_DB_PATH) as reader:
                    response = reader.country(ip)
                    country_name = response.country.name or response.country.iso_code
                    return f"{country_name} ({response.country.iso_code})"
            except:
                pass
        
        # 备用方案：使用简化的地理位置判断
        if is_foreign_ip(ip):
            return "国外IP (未知国家)"
        else:
            return "中国IP"
    except:
        return "未知位置"

def main():
    alerted = set()  # 记录已预警的(ip, url, minute)
    blocked_ips = set()  # 记录已拦截的IP
    
    # 初始化GeoIP2数据库
    geoip2_ready = init_geoip2_database()
    if geoip2_ready:
        logging.info("GeoIP2数据库初始化成功，将使用精确的地理位置检测")
    else:
        logging.info("将使用简化的中国IP网段检测")
    
    logging.info("启动Nginx日志监控，告警阈值=%s 次/分钟，拦截阈值=%s 次/分钟，窗口=%s秒，日志路径=%s", 
                THRESHOLD, BLOCK_THRESHOLD, INTERVAL, LOG_PATH)
    
    while True:
        now = datetime.now()
        window_start = now - timedelta(seconds=INTERVAL)
        logging.info("开始扫描窗口：%s - %s", window_start.strftime("%Y-%m-%d %H:%M:%S"), now.strftime("%Y-%m-%d %H:%M:%S"))
        
        ip_url_count = defaultdict(int)
        files_scanned = 0
        lines_read = 0
        lines_in_window = 0
        
        for log_file in glob.glob(LOG_PATH):
            files_scanned += 1
            with open(log_file, "r", encoding="utf-8", errors="ignore") as f:
                for line in f:
                    lines_read += 1
                    ip, dt, url = parse_log_line(line)
                    if ip and dt and url:
                        # 新增：排除指定URL前缀
                        if any(url.startswith(prefix) for prefix in EXCLUDE_URL_PREFIXES):
                            continue
                        # 新增：排除静态资源
                        if is_static_resource(url):
                            continue
                        if dt > window_start:
                            lines_in_window += 1
                            ip_url_count[(ip, url)] += 1
        
        logging.info(
            "扫描完成：文件=%d 行=%d 命中窗口=%d 唯一(IP,URL)=%d",
            files_scanned, lines_read, lines_in_window, len(ip_url_count)
        )
        
        alert_time = now.strftime("%Y-%m-%d %H:%M:%S")
        alerts_this_round = 0  
        blocks_this_round = 0
        
        for (ip, url), count in ip_url_count.items():
            minute_key = (ip, url, now.strftime("%Y%m%d%H%M"))
            
            if not is_public_ip(ip):
                continue
                
            # 检查是否需要拦截（只拦截国外IP且触发100次以上）
            should_block = False
            is_foreign = is_foreign_ip(ip)
            
            # 只有国外IP且触发100次以上才拦截
            if count >= BLOCK_THRESHOLD and is_foreign:
                should_block = True
                logging.warning("触发拦截条件：国外IP=%s 次数=%s 已超过拦截阈值", ip, count)
            elif count >= BLOCK_THRESHOLD and not is_foreign:
                logging.info("国内IP %s 触发%d次，但不在拦截范围内", ip, count)
            
            # 执行拦截（仅限国外IP）
            if should_block and ip not in blocked_ips:
                if block_ip_with_iptables(ip):
                    blocked_ips.add(ip)
                    blocks_this_round += 1
                    # 发送拦截告警
                    send_dingtalk_alert(ip, url, count, alert_time, is_blocked=True, is_foreign=is_foreign)
            
            # 检查是否需要普通告警（所有IP超过30次都告警，但不拦截）
            if count > THRESHOLD and minute_key not in alerted:
                send_dingtalk_alert(ip, url, count, alert_time, is_blocked=False, is_foreign=is_foreign)
                alerted.add(minute_key)
                alerts_this_round += 1
        
        if alerts_this_round == 0 and blocks_this_round == 0:
            logging.info("本轮无告警和拦截")
        else:
            logging.info("本轮告警：%d个，拦截：%d个", alerts_this_round, blocks_this_round)
            
        logging.info("休眠 %s 秒...", INTERVAL)
        time.sleep(INTERVAL)

if __name__ == "__main__":
    main()